using System.Collections.Generic;
using UnityEngine;

/// <summary>
/// Unity官方DrawMeshInstancedIndirect的示例.
/// https://docs.unity3d.com/ScriptReference/Graphics.DrawMeshInstancedIndirect.html
/// 追加Frustum culling裁剪掉视锥体之外的Instance再DrawMeshInstance. 一个最小化GPU Culling的实现. 这个Demo不支持对运动的Instance的Culling.
/// https://zhuanlan.zhihu.com/p/376801370
/// </summary>
public class FrustumCullingRunner : MonoBehaviour
{
    [Header("Instance")]
    public int instanceCount = 100000;

    public Mesh instanceMesh;
    public Material instanceMaterial;
    public int subMeshIndex = 0;
    private ComputeBuffer argsBuffer;
    private uint[] args = new uint[5] { 0, 0, 0, 0, 0 };

    [Header("Compute")]
    public ComputeShader compute;

    private ComputeBuffer localToWorldMatrixBuffer;
    private ComputeBuffer cullResult;
    private int kernel;
    private Camera mainCamera;

    private Vector4[] planes = new Vector4[6];

    private void OnEnable()
    {
        ClearBuffers();
        InitBuffers();
        InitInstanceMatrixs();
    }

    private void OnDisable()
    {
        ClearBuffers();
    }

    private void Update()
    {
        CullThenDraw();
    }
    /// <summary>
    /// 初始化所有的ComputeBuffer
    /// </summary>
    private void InitBuffers()
    {
        kernel = compute.FindKernel("ViewPortCulling");
        mainCamera = Camera.main;

        //CullResult是Matrix4x4, 16个float.
        cullResult = new ComputeBuffer(instanceCount, sizeof(float) * 16, ComputeBufferType.Append);

        //所有的Instance的local matrix, 因为Instance的transform不变化, 所以没有在update中更新.
        localToWorldMatrixBuffer = new ComputeBuffer(instanceCount, 16 * sizeof(float));

        //DrawMeshInstancedIndirect需要的参数Buffer. Mesh不变就不会改变, 所以只需要赋值1次
        argsBuffer = new ComputeBuffer(1, args.Length * sizeof(uint), ComputeBufferType.IndirectArguments);
        if (instanceMesh != null)
        {
            args[0] = (uint)instanceMesh.GetIndexCount(subMeshIndex);
            args[2] = (uint)instanceMesh.GetIndexStart(subMeshIndex);
            args[3] = (uint)instanceMesh.GetBaseVertex(subMeshIndex);
        }
        else
        {
            args[0] = args[1] = args[2] = args[3] = 0;
        }
        argsBuffer.SetData(args);

        // Ensure submesh index is in range
        if (instanceMesh != null)
        {
            subMeshIndex = Mathf.Clamp(subMeshIndex, 0, instanceMesh.subMeshCount - 1);
        }
    }

    /// <summary>
    /// 初始化Instance的Matrix(Position, Rotation, Scaling).
    /// </summary>
    private void InitInstanceMatrixs()
    {
        UnityEngine.Profiling.Profiler.BeginSample("InitInstanceMatrixs");

        //所有Instance的world matrix
        List<Matrix4x4> localToWorldMatrixs = new List<Matrix4x4>();
        for (int i = 0; i < instanceCount; i++)
        {
            float angle = Random.Range(0.0f, Mathf.PI * 2.0f);
            float distance = Random.Range(20.0f, 100.0f);
            float height = Random.Range(-2.0f, 2.0f);
            float size = Random.Range(0.05f, 0.25f);
            Vector4 position = new Vector4(Mathf.Sin(angle) * distance, height, Mathf.Cos(angle) * distance, size);
            localToWorldMatrixs.Add(Matrix4x4.TRS(position, Quaternion.identity, new Vector3(size, size, size)));
        }
        localToWorldMatrixBuffer.SetData(localToWorldMatrixs);

        UnityEngine.Profiling.Profiler.EndSample();
    }

    /// <summary>
    /// Cull之后Draw
    /// </summary>
    private void CullThenDraw()
    {
        //获取Frustum的6个面. 传入ComputShader
        mainCamera.GetFrustumPlaneNoneAllocate(planes);
        compute.SetVectorArray("planes", planes);

        //传入InstanceCount
        compute.SetInt("instanceCount", instanceCount);

        //传入所有Instance的
        compute.SetBuffer(kernel, "input", localToWorldMatrixBuffer);

        cullResult.SetCounterValue(0);
        compute.SetBuffer(kernel, "cullresult", cullResult);

        compute.Dispatch(kernel, 1 + (instanceCount / 640), 1, 1);

        //把cull结果传入shader.
        instanceMaterial.SetBuffer("positionBuffer", cullResult);

        //把cullResult.Count写入到argsBuffer中，只渲染Cull之后的个数.
        ComputeBuffer.CopyCount(cullResult, argsBuffer, sizeof(uint));

        //绘制Instance
        Graphics.DrawMeshInstancedIndirect(instanceMesh, subMeshIndex, instanceMaterial, 
            new Bounds(Vector3.zero, new Vector3(200.0f, 200.0f, 200.0f)),
            argsBuffer);
    }

    private void ClearBuffers()
    {
        localToWorldMatrixBuffer?.Release();
        localToWorldMatrixBuffer = null;

        cullResult?.Release();
        cullResult = null;

        argsBuffer?.Release();
        argsBuffer = null;

        subMeshIndex = 0;
    }
}